#define WIN32_LEAN_AND_MEAN

#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#include <time.h>

#include "rdistro.h"
#include "../tdp2/tdp2_hash.hpp"

namespace RDistro
{
	class DistributionSession
	{
	public:
		struct JobKey
		{
			UInt8 data[32];
			UInt16 jobTypeID;
		};

	private:
		clock_t _lastActivity;
		SOCKET _socket;
		bool _socketActive;
		bool _connectionActive;
		UInt8 _authenticationToken[32];
		UInt8 _loginKey[32];
		bool _loggedIn;
		int _connectionState;
		UInt8 _packetData[65536];

		UInt16 _ppProtocolID;
		UInt16 _ppSize;
		bool _ppWaiting;
		char _pendingFileName[4096];
		FILE *_activeFile;
		UInt8 _activeTimestamp;
		int _activeCompareResult;
		JobKey _activeJobKey;
		UInt32 _activeWorkFirstUnit;
		UInt32 _activeWorkNumUnits;

		void Initialize(const void *loginKey, size_t loginKeySize);
		bool FileTimestamp(const char *fname, UInt8 timestamp[8]);
		bool BlockOnState(int state);

	public:
		enum
		{
			CS_Disconnected,
			CS_Connecting,
			CS_WaitingForFeedback,
			CS_Idle,
			CS_SynchronizingFile,
			CS_UploadingFile,
			CS_DownloadingFile,
			CS_OperationFailed,
			CS_ConnectionDropped,
		};

		enum
		{
			JS_Finished,
			JS_Finishing,
			JS_Outstanding,
		};

		DistributionSession(const void *loginKey, size_t loginKeySize);
		explicit DistributionSession(const char *loginKey);
		~DistributionSession();

		void SendPacket(UInt16 protocolCode, const void *bytes, size_t numBytes);

		void ProcessPacket(int protocolID, const void *data, size_t size);
		void ProcessData();
		bool LoggedIn();
		bool ConnectionActive();
		int ConnectionState();
		void SetConnectionState(int state);
		
		bool UploadFile(const char *filename);
		bool SynchronizeFile(const char *filename, bool allowUpload = true, bool allowDownload = false, bool *updated = NULL);
		bool DownloadFile(const char *filename);
		bool RequestJob(JobKey &jobKey);
		bool RequestWork(const JobKey &jobKey, UInt32 &firstUnit, UInt32 &numUnits);
		void CreateJob(const JobKey &jobKey, UInt32 &numUnits, UInt8 &granularity);
	};
};

enum
{
	SSG_Read,
	SSG_Write,
	SSG_Except,
};

static bool selectsingle(SOCKET s, int grouping)
{
	TIMEVAL timeout;
	timeout.tv_usec = 0;
	timeout.tv_sec = 0;

	fd_set sockSet;
	sockSet.fd_array[0] = s;
	sockSet.fd_count = 0;

	fd_set *sets[3];
	sets[0] = sets[1] = sets[2] = NULL;

	sets[grouping] = &sockSet;

	return (select(1, sets[0], sets[1], sets[2], &timeout) != 0);
}

namespace RDistro
{
	void DistributionSession::Initialize(const void *loginKey, size_t loginKeySize)
	{
		_lastActivity = clock();

		_connectionState = CS_Disconnected;
		_socketActive = false;
		_connectionActive = false;
		_loggedIn = false;
		_ppWaiting = true;
		_activeFile = NULL;
		memset(_authenticationToken, 0, sizeof(_authenticationToken));
		TDP::Crypto::SHA256::SHA256Bytes(this->_loginKey, loginKey, loginKeySize);

		ADDRINFO hints;
		ADDRINFO *addrInfo;
		ADDRINFO *inetAddrInfo = NULL;
		memset(&hints, 0, sizeof(hints));
		getaddrinfo("localhost", "1791", &hints, &addrInfo);

		// Find INET
		for(ADDRINFO *subInfo = addrInfo;subInfo != NULL;subInfo = subInfo->ai_next)
		{
			if(subInfo->ai_family == AF_INET)
			{
				inetAddrInfo = subInfo;
				break;
			}
		}

		if(inetAddrInfo)
		{
			_socketActive = true;
			_socket = socket(inetAddrInfo->ai_family, inetAddrInfo->ai_socktype, inetAddrInfo->ai_protocol);

			{
				u_long mode = 1;
				ioctlsocket(_socket, FIONBIO, &mode);
			}

			int conRV = connect(_socket, inetAddrInfo->ai_addr, static_cast<int>(inetAddrInfo->ai_addrlen));
			if(conRV != SOCKET_ERROR || WSAGetLastError() == WSAEWOULDBLOCK)
			{
				_connectionState = CS_Connecting;
			}
		}

		freeaddrinfo(addrInfo);
	}

	bool DistributionSession::FileTimestamp(const char *fname, UInt8 timestamp[8])
	{
		HANDLE f = CreateFile(fname, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
		if(f == INVALID_HANDLE_VALUE)
			return false;

		FILETIME modTime;
		GetFileTime(f, NULL, NULL, &modTime);
		CloseHandle(f);

		timestamp[0] = static_cast<UInt8>(modTime.dwHighDateTime >> 24);
		timestamp[1] = static_cast<UInt8>(modTime.dwHighDateTime >> 16);
		timestamp[2] = static_cast<UInt8>(modTime.dwHighDateTime >> 8);
		timestamp[3] = static_cast<UInt8>(modTime.dwHighDateTime);
		timestamp[4] = static_cast<UInt8>(modTime.dwLowDateTime >> 24);
		timestamp[5] = static_cast<UInt8>(modTime.dwLowDateTime >> 16);
		timestamp[6] = static_cast<UInt8>(modTime.dwLowDateTime >> 8);
		timestamp[7] = static_cast<UInt8>(modTime.dwLowDateTime);
		return true;
	}

	DistributionSession::DistributionSession(const void *loginKey, size_t loginKeySize)
	{
		Initialize(loginKey, loginKeySize);
	}
	
	DistributionSession::DistributionSession(const char *loginKey)
	{
		Initialize(loginKey, strlen(loginKey));
	}
	
	DistributionSession::~DistributionSession()
	{
		if(_socketActive)
		{
			closesocket(_socket);
			_socketActive = false;
		}
		if(_activeFile)
		{
			fclose(_activeFile);
			_activeFile = NULL;
		}
	}

	void DistributionSession::SendPacket(UInt16 protocolCode, const void *bytes, size_t numBytes)
	{
		UInt8 hash[32];
		UInt8 byteBuf[32+32+4+65535+32];

		memcpy(byteBuf, _authenticationToken, 32);
		memcpy(byteBuf + 32, _loginKey, 32);
		byteBuf[64] = static_cast<UInt8>(numBytes);
		byteBuf[65] = static_cast<UInt8>(numBytes >> 8);
		byteBuf[66] = static_cast<UInt8>(protocolCode);
		byteBuf[67] = static_cast<UInt8>(protocolCode >> 8);
		memcpy(byteBuf + 68, bytes, numBytes);
		TDP::Crypto::SHA256::SHA256Bytes(hash, byteBuf, numBytes + 68);
		memcpy(byteBuf + 32, hash, 32);

		int rv = send(_socket, reinterpret_cast<const char *>(byteBuf + 32), static_cast<int>(32+4+numBytes), 0);
	}

	int DistributionSession::ConnectionState()
	{
		return _connectionState;
	}

	void DistributionSession::SetConnectionState(int state)
	{
		_connectionState = state;
	}

	bool DistributionSession::ConnectionActive()
	{
		return _connectionState != CS_ConnectionDropped;
	}

	bool DistributionSession::LoggedIn()
	{
		return _loggedIn;
	}

	void DistributionSession::ProcessPacket(int protocolID, const void *data, size_t size)
	{
		switch(protocolID)
		{
		case CL_PING:
			{
				SendPacket(SV_KEEPALIVE, NULL, 0);
			}
			break;
		case CL_SET_AUTH_TOKEN:
			{
				if(size == 32)
				{
					memcpy(this->_authenticationToken, data, 32);
					this->_loggedIn = true;
				}
			}
			break;
		case CL_UPLOAD_FAILED:
			{
				if(_activeFile)
				{
					fclose(_activeFile);
					_activeFile = NULL;
				}
				_connectionState = CS_OperationFailed;
			}
			break;
		case CL_UPLOAD_MORE:
			{
				UInt8 data[16000];
				size_t nBytes = fread(data, 1, sizeof(data), _activeFile);

				if(nBytes)
					SendPacket(SV_UPLOAD_FILE_DATA, data, nBytes);
				else
				{
					fclose(_activeFile);
					_activeFile = NULL;

					UInt8 timestamp[8];

					if(!FileTimestamp(_pendingFileName, timestamp))
					{
						_connectionState = CS_OperationFailed;
						break;
					}
					_connectionState = CS_Idle;

					SendPacket(SV_UPLOAD_FILE_FINISH, timestamp, 8);
				}
			}
			break;
		case CL_TIME_CHECK_FILE_MATCHES:
		case CL_TIME_CHECK_FILE_LOCAL_IS_NEWER:
		case CL_TIME_CHECK_FILE_SERVER_IS_NEWER:
			{
				_activeCompareResult = protocolID;
				_connectionState = CS_Idle;
			}
			break;
		case CL_FILE_DOWNLOAD_INTERRUPTED:
		case CL_FILE_DOES_NOT_EXIST:
		case CL_WORK_EXHAUSTED:
			{
				_connectionState = CS_OperationFailed;
			}
			break;
		case CL_DOWNLOAD_FILE:
			{
				_connectionState = CS_DownloadingFile;
				_activeFile = fopen(_pendingFileName, "wb");
				SendPacket(SV_DOWNLOAD_MORE, NULL, 0);
			}
			break;
		case CL_DOWNLOAD_FILE_DATA:
			{
				fwrite(data, size, 1, _activeFile);
				SendPacket(SV_DOWNLOAD_MORE, NULL, 0);
			}
			break;
		case CL_DOWNLOAD_FILE_FINISH:
			{
				_connectionState = CS_Idle;
				fclose(_activeFile);

				const UInt8 *timestamp = reinterpret_cast<const UInt8 *>(data);

				FILETIME convertedStamp;
				convertedStamp.dwHighDateTime = (timestamp[0] << 24) | (timestamp[1] << 16) | (timestamp[2] << 8) | (timestamp[3]);
				convertedStamp.dwLowDateTime = (timestamp[4] << 24) | (timestamp[5] << 16) | (timestamp[6] << 8) | (timestamp[7]);

				HANDLE f = CreateFile(_pendingFileName, FILE_WRITE_ATTRIBUTES, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
				if(f != INVALID_HANDLE_VALUE)
				{
					BOOL result = SetFileTime(f, NULL, NULL, &convertedStamp);
					CloseHandle(f);
				}
			}
			break;
		case CL_WORK_ON_JOB:
			{
				const UInt8 *bdata = reinterpret_cast<const UInt8 *>(data);
				memcpy(_activeJobKey.data, data, 32);
				_activeJobKey.jobTypeID = (bdata[32] | (bdata[33] << 8));
				_connectionState = CS_Idle;
			}
			break;
		case CL_WORK:
			{
				const UInt8 *bdata = reinterpret_cast<const UInt8 *>(data);
				_activeWorkFirstUnit = (bdata[0] | (bdata[1] << 8) | (bdata[2] << 16) | (bdata[3] << 24));
				_activeWorkNumUnits = (bdata[4] | (bdata[5] << 8) | (bdata[6] << 16) | (bdata[7] << 24));
				_connectionState = CS_Idle;
			}
			break;
		};
	}

	void DistributionSession::ProcessData()
	{
		if(_connectionState == CS_ConnectionDropped)
			return;

		if(clock() - _lastActivity > 20*CLOCKS_PER_SEC)
		{
			_connectionState = CS_ConnectionDropped;
			_loggedIn = false;
			closesocket(_socket);
			return;
		}


		int bytesNeeded = 4;
		if(!_ppWaiting)
			bytesNeeded = this->_ppSize;

		int rv = recv(_socket, reinterpret_cast<char *>(_packetData), bytesNeeded, 0);
		if(rv == SOCKET_ERROR)
		{
			if(WSAGetLastError() != WSAEWOULDBLOCK)
			{
				_connectionState = CS_ConnectionDropped;
				_loggedIn = false;
				closesocket(_socket);
			}
		}
		else
		{
			if(rv != bytesNeeded)
			{
				_connectionState = CS_ConnectionDropped;
				_loggedIn = false;
				closesocket(_socket);
			}
			else
			{
				if(_ppWaiting)
				{
					// Received the header
					_ppSize = _packetData[0] | (_packetData[1] << 8);
					_ppProtocolID = _packetData[2] | (_packetData[3] << 8);
					
					if(_ppSize == 0)
						ProcessPacket(_ppProtocolID, NULL, 0);
					else
						_ppWaiting = false;
				}
				else
				{
					ProcessPacket(_ppProtocolID, _packetData, _ppSize);
					_ppWaiting = true;
				}
			}
		}
	}
	
	bool DistributionSession::SynchronizeFile(const char *filename, bool allowUpload, bool allowDownload, bool *updated)
	{
		_lastActivity = clock();

		if(updated)
			*updated = false;

		UInt8 resultRequestPacket[4096];
		if(!FileTimestamp(filename, resultRequestPacket))
		{
			if(!allowDownload)
				return true;
			if(updated)
				*updated = true;
			return DownloadFile(filename);
		}

		strcpy(_pendingFileName, filename);
		strcpy(reinterpret_cast<char *>(resultRequestPacket + 8), filename);
		
		SetConnectionState(CS_SynchronizingFile);

		SendPacket(SV_TIME_CHECK_FILE, resultRequestPacket, strlen(filename) + 8);

		if(!BlockOnState(CS_SynchronizingFile))
			return false;

		printf("Sync compare result: %i\n", _activeCompareResult);

		if(this->_activeCompareResult == CL_TIME_CHECK_FILE_LOCAL_IS_NEWER)
		{
			if(!allowUpload)
				return true;
			return UploadFile(filename);
		}
		if(this->_activeCompareResult == CL_TIME_CHECK_FILE_SERVER_IS_NEWER)
		{
			if(!allowDownload)
				return true;
			if(updated)
				*updated = true;
			return DownloadFile(filename);
		}
		return true;
	}

	bool DistributionSession::BlockOnState(int state)
	{
		while(_connectionState == state)
		{
			Sleep(1);
			ProcessData();
		}
		
		if(_connectionState == CS_OperationFailed)
		{
			_connectionState = CS_Idle;
			return false;
		}
		return (_connectionState == CS_Idle);
	}

	bool DistributionSession::UploadFile(const char *filename)
	{
		printf("Uploading %s\n", filename);

		_lastActivity = clock();
		_activeFile = fopen(filename, "rb");
		if(!_activeFile)
			return false;

		strcpy(_pendingFileName, filename);
		
		SetConnectionState(CS_UploadingFile);

		SendPacket(SV_UPLOAD_FILE, filename, strlen(filename));
		
		return BlockOnState(CS_UploadingFile);
	}

	bool DistributionSession::DownloadFile(const char *filename)
	{
		printf("Downloading %s\n", filename);

		_lastActivity = clock();

		strcpy(_pendingFileName, filename);

		SetConnectionState(CS_DownloadingFile);

		SendPacket(SV_DOWNLOAD_FILE, filename, strlen(filename));

		return BlockOnState(CS_DownloadingFile);
	}

	bool DistributionSession::RequestJob(JobKey &jobKey)
	{
		_lastActivity = clock();
		
		SetConnectionState(CS_WaitingForFeedback);
		
		SendPacket(SV_REQUEST_JOB, NULL, 0);
		
		if(!BlockOnState(CS_WaitingForFeedback))
			return false;

		jobKey = _activeJobKey;

		return true;
	}

	bool DistributionSession::RequestWork(const JobKey &jobKey, UInt32 &firstUnit, UInt32 &numUnits)
	{
		_lastActivity = clock();

		SetConnectionState(CS_WaitingForFeedback);

		SendPacket(SV_REQUEST_MORE_WORK, jobKey.data, 32);
		
		if(!BlockOnState(CS_WaitingForFeedback))
			return false;

		firstUnit = _activeWorkFirstUnit;
		numUnits = _activeWorkNumUnits;

		return true;
	}

	void DistributionSession::CreateJob(const JobKey &jobKey, UInt32 numUnits, UInt8 granularity, const char *target)
	{
		UInt8 request[4096];
		memcpy(request, jobKey.data, 32);
		request[32] = static_cast<UInt8>(jobKey.jobTypeID);
		request[33] = static_cast<UInt8>(jobKey.jobTypeID >> 8);
		request[34] = static_cast<UInt8>(numUnits);
		request[35] = static_cast<UInt8>(numUnits >> 8);
		request[36] = static_cast<UInt8>(numUnits >> 16);
		request[37] = static_cast<UInt8>(numUnits >> 24);
		request[38] = granularity;
		strcpy(reinterpret_cast<char *>(request + 39, target);

		SendPacket(SV_CREATE_JOB, request, 39 + strlen(target));
	}
};
